import os
import math

import cv2
import numpy as np
from mayavi import mlab
from scipy.optimize import leastsq, lsq_linear  # 优化和拟合函数库scipy.optimize
from tvtk.api import tvtk
from skimage import io,data
from read_data import *

np.set_printoptions(threshold=np.inf)   #设置array全部输出
np.seterr(divide='ignore', invalid='ignore')


def LPpara(Ip,k,d,E0):
    #E0 = read(sheet_index=2, m0=k-1, m=k)[1]
    I0 = np.array([[0],[0],[-1]])
    Ep = (E0*((np.vdot(I0,Ip))**(g-1)))/(d**2)
    #print('E0:',E0,' np:',np.vdot(I0,Ip),' d:',d,' g:',g)
    Lp = Ep*Ip
    return Lp
#E0,I0,d,g

def CoorTransU2x(u,v,u0,v0,dx,dy):
    x=(u-u0)*dx
    y=(v-v0)*dy
    return x,y

def InitSur(x,y,f):
    return x/f,y/f,-1
    #return 12.5294*x/f,12.5294*y/f,-12.5294

def LightDirection(k,xb,yb,zb,Lx,Ly,Lz):
    #Lx, Ly, Lz = read(sheet_index=1,m0=k,m=k+1,n=4)
    #print(Lx,Ly,Lz)
    #return np.array([[-xb+Lx],[-yb+Ly],[-zb+Lz]])/math.sqrt((xb-Lx)**2+(yb-Ly)**2+(zb-Lz)**2)
    return np.array([[xb-Lx],[yb-Ly],[zb-Lz]])/math.sqrt((xb-Lx)**2+(yb-Ly)**2+(zb-Lz)**2)


def OneImage(k,imgdir):

    img = cv2.imread(imgdir)
    img_gray = cv2.cvtColor(img,cv2.COLOR_RGB2GRAY)
    #cv2.imshow('img_gray',img_gray)
    im = np.array(img_gray)
    #将灰度图转化为数组im，值为图像亮度
    length,width = im.shape
    #图像尺寸 length=600,width=800
    E0 = read(sheet_index=2, m0=k-1, m=k)[0]
    #读E0
    mIIp = np.zeros((600,600))
    mLp = np.zeros((600,600,3))
    '''
    mIIp = np.zeros((600,800))
    mLp = np.zeros((600,800,3))
    '''
    #IIp、Lp矩阵
    Lx, Ly, Lz = read(sheet_index=1, m0=k-1, m=k, n=4)
    #光源坐标
    for u in range(140,740):
    #width
        for v in range(600):
    #length
            IIp = im[v,u]//10
            #每一点反射亮度
            xi0, yi0 = CoorTransU2x(u, v, u0, v0, dx, dy)
            xi, yi, zi = InitSur(xi0, yi0, f)
            #初始化坐标 xi=xi0/f, yi=yi0/f, z=1
            d = math.sqrt((xi-Lx)**2+(yi-Ly)**2+(zi-Lz)**2)
            #求d
            Ip = LightDirection(k, xi, yi, zi, Lx, Ly, Lz)
            Lp = LPpara(Ip,k,d,E0)
            #求Lp
            #print(u,v,' IIp:',IIp,' Lp:',Lp)
            mIIp[v,u-140] = IIp
            mLp[v,u-140] = list(Lp)
            #将每一点的IIp、Lp存入矩阵mIIp、mLp
    return mIIp,mLp,length,width
#Ip:入射光线；IIp:反射亮度
# u0,y0,dx,dy,k


def func(p,x,y,z):
    a,b,c = p
    return a*x+b*y+c*z

def error(p,x,y,z,m):
    return func(p,x,y,z)-m

def normal(xi,yi,zi,mi):
    p0 = [1,1,1]
    Para = leastsq(error,p0,args=(xi, yi, zi, mi))
    a,b,c = Para[0]
    return a,b,c
#最小二乘法


def NormalDirection(imgdir):
    allmIIp = np.zeros((8,600,600))
    allmLp = np.zeros((8,600,600,3))
    '''
    allmIIp = np.zeros((8,600,800))
    allmLp = np.zeros((8,600,800,3))
    '''
    #初始化8副图的总矩阵
    for k in range(1,9):
        imgdir = imgdir.replace(str(k-1)+'.jpg',str(k)+'.jpg')
        mIIp,mLp,length,width = OneImage(k,imgdir)
        allmIIp[k-1] = mIIp
        allmLp[k-1] = mLp
        #print(k)
        #将每幅图的的mIIp、mLp存入allmIIp、allmLp
    #print(allmIIp)
    #print()
    #print(allmLp)
    #print()
    allNp = np.zeros((600,600,3))
    xi = np.zeros((600,600,8))
    yi = np.zeros((600,600,8))
    zi = np.zeros((600,600,8))
    mi = np.zeros((600,600,8))
    #初始化xi、yi、zi、mi

    for i in range(600):
        #800
        for j in range(600):
        #600
            for k in range(8):
                xi[j,i][k] = allmLp[k][j,i][0]
                yi[j,i][k] = allmLp[k][j,i][1]
                zi[j,i][k] = allmLp[k][j,i][2]
                mi[j,i][k] = allmIIp[k][j,i]
                #将每个点上，8张图的数据存入xi、yi、zi、mi
            allNp[j,i] = normal(xi[j,i],yi[j,i],zi[j,i],mi[j,i])
            #对每个点进行最小二乘法求解，求得各点法向量存入allNp

    return allmIIp, allmLp, allNp
    

#FrankotChellappa算法,参数p,q为梯度
def frankotchellappa(p, q):
    M, N = p.shape
    #对于非周期对象的梯度进行扩展
    p = np.concatenate((np.hstack((p, -p[:,::-1])),np.hstack((p[::-1,:], -p[::-1,::-1]))))
    q = np.concatenate((np.hstack((q, q[:,::-1])),np.hstack((-q[::-1,:], -q[::-1,::-1]))))
    M *= 2
    N *= 2
    #计算频率指数
    half_M = (M - 1) / 2
    half_N = (N - 1) / 2
    u, v = np.meshgrid(np.r_[-np.ceil(half_N):np.floor(half_N)+1],
                        np.r_[-np.ceil(half_M):np.floor(half_M)+1])
    u = np.fft.ifftshift(u)
    v = np.fft.ifftshift(v)
    #计算梯度p、q的离散傅里变换
    Fp = np.fft.fft2(p)
    Fq = np.fft.fft2(q)
    #计算高度Z的傅里叶变换
    Fz = -1j / (2 * np.pi) * (u * Fp / N + v * Fq / M) / ((u / N)** 2 + (v / M)** 2)
    #Z傅里叶变换常量置零
    Fz[0,0] = 0
    #高度Z傅里叶反变换
    Z = np.real(np.fft.ifft2(Fz))
    #恢复非周期对象的高度Z
    result = Z[0:M // 2, 0:N // 2]
    return result


#利用linear least-squares进行三次多项式拟合曲面，参数Z为待拟合数据（二维数组）
def polynomial_fit(Z):
    '''
    fitresult(x,y) = p00 + p10*x + p01*y + p20*x^2 + p11*x*y + p02*y^2 
                    + p30*x^3 + p21*x^2*y + p12*x*y^2 + p03*y^3
    A = [1, x, y, x ^ 2, x * y, y ^ 2, x ^ 3, x ^ 2 * y, x * y ^ 2, y ^ 3]
    x = [p00, p10, p01, p20, p11, p02, p30, p21, p12, p03]
    b = [z = f(x, y)]
    Ax = b -> x
    '''
    #拟合数据预处理，将需要拟合的数据转为列向量
    n = Z.shape[0]
    X, Y = np.meshgrid(np.r_[0:n], np.r_[0:n])
    xdata = X.reshape(-1, 1, order='F')
    ydata = Y.reshape(-1, 1, order='F')
    zdata = Z.reshape(-1, order='F')
    #构建矩阵A，b
    x0 = np.ones(xdata.shape, dtype=xdata.dtype)
    x1 = xdata
    y1 = ydata
    x2 = xdata ** 2
    xy = xdata * ydata
    y2 = ydata * ydata
    x3 = xdata ** 3
    x2y = xdata **2 * ydata
    xy2 = xdata * ydata ** 2
    y3 = ydata ** 3
    A = np.c_[x0, x1, y1, x2, xy, y2, x3, x2y, xy2, y3]
    b = zdata
    #利用linear least-squares求解多项式系数
    res = lsq_linear(A, b)
    p = res.x
    #输出拟合结果
    fitresult = np.dot(A, p).reshape(n, -1, order='F')
    return fitresult



def reconstruct_z(file_path):
    global u0, v0, f, dx, dy, g
    u0, v0, f, dx, dy = read(m=5)
    u0 = 400
    v0 = 300
    dx = (dx * 2592)/800
    dy = (dy * 2592)/800
    g = 3
    imgdir = file_path + os.sep + 'images' + os.sep + '0.jpg'
    
    IIp,Lp,Np = NormalDirection(imgdir)
    '''
    print(type(Np))
    print(Np)
    '''
    #print(Np.shape)
    #print()

    PQ = np.zeros((600,600,2))
    #    PQ = np.zeros((600,800,2))
    for i in range(600):
        #600
        for j in range(600):
            #800
            PQ[599-i,j][0] = Np[i,j][0]/Np[i,j][2]
            PQ[599-i,j][1] = Np[i,j][1]/Np[i,j][2]

            Np[i,j][0] = Np[i,j][0]/Np[i,j][2]
            Np[i,j][1] = Np[i,j][1]/Np[i,j][2]
            Np[i,j][2] = Np[i,j][2]/Np[i,j][2]
    '''
    txtname = './temp/NormalVector.txt'
    f = open(txtname, 'w')
    for i in range(600):
        for j in range(600):
            a = str(PQ[i,j]).strip('[')
            a = a.strip(']')
            f.writelines(a+' ')
        f.writelines('\n')
    f.close()
    '''
    img=io.imread(file_path + os.sep + 'images' + os.sep + '0.jpg')
    roi=img[0:600,140:740,:]
    #io.imshow(roi)
    #plt.show()
    io.imsave(file_path + os.sep + 'display.jpg', roi)
    roi = np.rot90(roi)

    io.imsave(file_path + os.sep + 'texture.jpg', roi)



    '''
    img = cv2.imread(file_path + os.sep + 'images' + os.sep + '0.jpg')
    roi = img[0:600, 140:740,:]
    cv2.imwrite(file_path + os.sep + 'display.jpg', roi)
    #roi = np.rot90(roi)
    #roi = np.rot90(roi)
    roi = np.flipud(roi)
    cv2.imwrite(file_path + os.sep + 'texture.jpg', roi)
    '''
    np.save('./temp/PQ.npy', PQ)

    p = PQ[:,:, 0]
    q = PQ[:,:, 1]
    
    Z = frankotchellappa(p, q)

    Z=np.flipud(Z)

    fitresult = polynomial_fit(Z)

    Z = Z - fitresult
    Z = Z * 1.33
    

    np.save(file_path + os.sep + 'Z.npy', Z)

    #plot_surface(Z)

if __name__ == '__main__':
    reconstruct_z('.')
